# -*- coding: utf-8 -*-
# @author: HRUN
import copy
import importlib
import json
import subprocess
from collections import OrderedDict
import re
import os
from numbers import Number
from locust import FastHttpUser, task, TaskSet, between, events, LoadTestShape, HttpUser, constant
from locust.env import Environment
from gevent._semaphore import Semaphore

from requests_toolbelt import MultipartEncoder

global ENV, global_func, global_func_file


try:
    global_func = importlib.import_module('global_func')
except ModuleNotFoundError:
    from apitestengine.core import tools as global_func

os.environ.setdefault("DJANGO_SETTINGS_MODULE", "primaryApp.settings.dev")
import django
django.setup()




# 集合点逻辑实现
from projects.models import TestEnv
from performanceengine.params import load_data
all_locusts_spawned = Semaphore()
all_locusts_spawned.acquire()  # 阻塞线程

def on_hatch_complete(user_count, **kwargs):
    """
    Select_task类的钩子方法
    """
    # 创建钩子方法
    all_locusts_spawned.release()
events.spawning_complete.add_listener(on_hatch_complete)

n = 0

def get_env_config(env_id, debug=True):
    """获取测试环境的配置"""
    try :
        env = TestEnv.objects.get(pk=env_id)
        var = {**env.global_variable, **env.debug_global_variable} if debug else env.global_variable
        # 环境变量
        ENV = {
            **var,  # 因为debug模式下var中有临时的调试变量，解包时不能写到下面，否则会覆盖真正的host和headers
            'host': env.host,
            'headers': env.headers,
        }
        config = {
            'ENV': ENV,
            'DB': env.db,
            'global_func': env.global_func
        }
        return config
    except Exception :
        return {}



class BaseEnv(dict):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def __setattr__(self, key, value):
        super().__setitem__(key, value)

    def __getattr__(self, item):
        return super().__getitem__(item)
ENV = BaseEnv()
# 不同场景存在不同环境
class SubEnvironment(Environment):
    def __init__(self):
        super().__init__()
        self.url = None

class BaseStepDispose():
    @classmethod
    def setUpClass(cls) -> None:
        cls.env = ENV

    def save_global_variable(self, name, value):
        ENV[name] = value

    def api_perform(self, step):

        try:
            data = copy.deepcopy(step).get('content', {})
            # 执行前置脚本
            self.__run_setup_script(data)
            # 发送请求
            response = self.__send_locustRequest(data)
            # 执行后置脚本
            self.__run_teardown_script(response)
        except Exception as e:
            raise Exception(f"{e}")




    def __send_locustRequest(self, data):
        method = data.get('method', 'POST').upper()  # 统一转大写
        request_info = self.construction_headers_data(data)
        # 统一处理所有请求类型
        valid_methods = {'GET', 'POST', 'PUT', 'DELETE', 'HEAD', 'PATCH'}
        if method not in valid_methods:
            raise ValueError(f"Invalid HTTP method: {method}")

        try:
            # 动态调用请求方法（核心优化点）
            with self.client.request(
                    method=method,
                    **request_info,
                    catch_response=True
            ) as response:
                if 200 <= response.status_code < 300:
                    response.success()
                    # print(f"传参：{request_info}", f"成功响应: {response.text}")
                    return response
                else:
                    response.failure(f"状态码异常: {response.status_code}")
                    return None

        except Exception as e:
            # 统一异常处理
            print(f"请求失败: {str(e)}")
            return None

    def construction_headers_data(self, data):
        """处理请求数据"""
        # 获取请求头
        if ENV.get('headers'):
            data['headers'] = {**ENV.get('headers'), **data.get('headers')}
        # 替换用例数据中的变量
        for k, v in list(data.items()):
            if k in ['interface', "headers", 'request', 'file']:
                # 替换变量
                v = self.__parser_variable(v)
                data[k] = v
        # files字段文件上传处理的处理
        files = data.get('file')
        if files:
            if isinstance(files, dict):
                file_data = files.items()
            else:
                file_data = files
            field = []
            for name, file_info in file_data:
                # 判断是否时文件上传(获取文件类型和文件名)
                if len(file_info) == 3 and os.path.isfile(file_info[1]):
                    field.append([name, (file_info[0], open(file_info[1], 'rb'), file_info[2])])
                else:
                    field.append([name, file_info])
            form_data = MultipartEncoder(fields=field)
            data['headers']["Content-Type"] = form_data.content_type
            data['request']['data'] = form_data
            data['files'] = None
        else:
            pass
        # 组织requests 发送请求所需要的参数格式
        request_params = {}
        # requests请求所需的所有字段
        params_fields = ['url', 'method', 'params', 'data', 'json', 'files', 'headers', 'cookies', 'auth', 'timeout',
                         'allow_redirects', 'proxies', 'hooks', 'stream', 'verify', 'cert']
        for k, v in data['request'].items():
            if k in params_fields:
                request_params[k] = v

        request_params['url'] = data.get('url')
        request_params['name'] = data.get('name')
        # 请求头
        request_params['headers'] = data['headers']
        return request_params

    def __run_teardown_script(self, response):
        """执行后置脚本"""
        self._hook_gen.send(response)
        delattr(self, '_hook_gen')
    def __run_setup_script(self, data):
        """执行前置脚本"""
        self._hook_gen = self.__run_script(data)
        next(self._hook_gen)
    def __run_script(test, data):
        # env = test.env
        setup_script = data.get('setup_script')
        if setup_script:
            try:
                exec(setup_script)
            except Exception as e:
                delattr(test, '_hook_gen')
                raise
        response = yield
        teardown_script = data.get('teardown_script')
        if teardown_script:
            try:
                exec(teardown_script)
            except AssertionError as e:
                raise e
            except Exception as e:
                raise
        yield

    def __parser_variable(self, data):
        """替换变量"""
        pattern = r'\{{(.+?)}}'
        old_data = data
        if isinstance(data, OrderedDict):
            data = dict(data)
        data = str(data)

        while re.search(pattern, data):
            res2 = re.search(pattern, data)
            item = res2.group()
            attr = res2.group(1)
            value = ENV.get(attr)
            if value is None:
                raise ValueError('变量引用错误：\n{}\n中的变量{},在当前运行环境中未找到'.format(
                    json.dumps(old_data, ensure_ascii=False, indent=2), attr)
                )
            if isinstance(value, Number):
                s = data.find(item)
                dd = data[s:s + len(item)]
                data = data.replace(dd, str(value))
            elif isinstance(value, str) and "'" in value:
                data = data.replace(item, value.replace("'", '"'))
            else:
                data = data.replace(item, str(value))
        return eval(data)
    def if_perform(self, case):
        pass

    def script_perform(self, case):
        pass


class GenerateTask():
    def __init__(self, control='20'):
        self.control = control
    def create_taskSet_class(self, item: list) -> list:
            task_class = []
            for idx, scence in enumerate(item):
                cls_name = scence.get('name') or 'Demo' + f'_{idx + 1}'
                weight = int(scence.get('weight', 1))
                task_funcs = {}
                if scence.get('steps'):
                    for step_idx, step in enumerate(scence['steps']):
                        func_name = self.create_testTask_name(idx, step_idx, step)
                        test_method = self.create_testTask_func(step)
                        task_funcs[func_name] = test_method
                task_funcs['on_start'] = self.on_start
                task_funcs['on_stop'] = self.on_stop

                task_funcs['weight'] = weight
                cls = type(cls_name, (TaskSet,BaseStepDispose,), task_funcs)
                task_class.append(cls)

            return task_class

    def on_start(self):
        '''场景集初始化执行方法'''
        global n
        if self.control == '10':
            n += 1
            print("第%s个虚拟用户开始启动" % n)
            all_locusts_spawned.wait()  # 同步锁等待
    def on_stop(self):
        """场景集结束执行方法"""
        pass
    def create_testTask_name(self, case_index, step_index, step):
        """生成唯一的测试方法名，避免重复"""
        step_title = step.get('title', 'UnknownTitle').replace(' ', '_')
        testTask_name = f"testTask_{case_index + 1}_{step_index + 1}_{step_title}"
        return testTask_name
    def create_testTask_func(self, step):
        """创建测试方法，使用函数工厂来动态传递参数"""
        type = step.get('type')
        def test_method(self):
            if type == 'api':
                return self.api_perform(step)
            elif type == 'script':
                return self.script_perform(step)
            elif type == 'if':
                return self.if_perform(step)
            else:
                raise ValueError(f"不支持的步骤类型:{type}")
        if type == 'api':
            weight = step.get('weight', 1)
            return task(weight)(test_method)

        return test_method


class Config():
    def __init__(self, config,default_url = ""):
        self.config = config
        self.default_url = default_url
        self.write_config()

    def write_config(self):
        """写入配置文件"""
        # config_lines = ["headless = true\n"]
        config_lines = ["headless = false\n", f"host = {self.default_url}\n"]
        if isinstance(self.config, dict):
            time_unit = None
            pressureMode = None
            for key, value in self.config.items():
                if key == 'timeUnit':
                    time_unit = value
                elif key == 'pressureMode':
                    pressureMode = value
                elif key == 'pressureConfig':
                    if pressureMode == '10':
                        config_lines.extend(self.process_pressure_config(value, time_unit))

                elif key == 'serverArray':
                    self.server(value)

                elif key == 'logMode':
                    self.run_log(value)

        # 写入配置文件
        try:
            with open("locust.conf", "w") as config_file:
                config_file.writelines(config_lines)
        except IOError as e:
            print(f"Error writing to file: {e}")
            # 错误处理重试处理

    def process_pressure_config(self, pressure_config, time_unit) -> list:
        """处理 pressureConfig 配置"""
        lines = []
        for key, value in pressure_config.items():
            if key == 'lastLong' and time_unit:
                lines.append(f"run-time = {value}{time_unit}\n")
            elif key == 'concurrencyNumber':
                lines.append(f"users = {value}\n")
            elif key == 'concurrencyStep':
                lines.append(f"spawn-rate = {value}\n")


        return lines

    def server(self, server_array: list):
        """启动locust性能压测的服务器配置"""
        pass
    def run_log(self, log_mode: int):
        """
        运行过程中日志输出粒度：
        0：关闭
        10：开启-全部日志
        20：开启-仅成功日志
        30：开启-仅失败日志
        """
        pass

    def config_file(self) -> dict:
        """配置项返回给其他模块使用"""
        config = self.config
        filter_config = {key: value for key, value in config.items() if key not in
                         ['id', 'serverNames','creator','create_time','modifier',
                          'update_time', 'name',  'isSetting', 'project',
                          'task','serverArray', 'resource', ]}

        stages = []
        if filter_config.get('pressureMode', None) == '20':
            config = filter_config.get('pressureConfig')
            ladders = config.get('ladders')
            sum_duration = 0
            for stage in ladders:
                lastLong = stage.get('lastLong',0)
                duration = self.timeUnit(filter_config.get('timeUnit', 's'), lastLong)
                sum_duration += duration
                stages.append({
                    "duration": sum_duration,
                    "users": int(stage.get('concurrencyNumber',0)),
                    "spawn_rate": int(stage.get('concurrencyStep', 0)),
                })
            filter_config['stages'] = stages
            del filter_config['pressureConfig']

        # print(filter_config)
        return  filter_config


    def timeUnit(self, time_unit, value):
        """时间单位转换"""
        value = int(value)
        if time_unit == 's':
            return value
        elif time_unit == 'm':
            return value * 60
        elif time_unit == 'h':
            return value * 3600
        else:
            return value






item = load_data()
env_config = get_env_config(item.get('env'))
ENV = {**env_config.get('ENV', {})}
global_func_file = ENV.get('global_func', b'')
if global_func_file:
    exec(global_func_file, global_func.__dict__)

conf = Config(item.get('presetting', {}),env_config.get('ENV').get('host'))
runConf = conf.config_file()
# 创建 GenerateTask 实例
generator = GenerateTask(runConf.get('control'))

cls_list = generator.create_taskSet_class(item.get('scenes'))



class CreateKitClass(FastHttpUser):
    tasks = cls_list
    if runConf.get('thinkTimeType') == '10':
        wait_time = constant(runConf.get('thinkTime')[0])
    else:
        wait_time = between(runConf.get('thinkTime')[0], runConf.get('thinkTime')[1])
    @events.test_start.add_listener
    def on_test_start(environment, **kwargs):
        print("开始执行任务集初始化动作")
    @events.test_stop.add_listener
    def on_test_stop(environment, **kwargs):
        print("结束执行任务集初始化动作")


if runConf.get('pressureMode') == '20':
    class StagesShapeWithCustomUsers(LoadTestShape):
        """
        A simply load test shape class that has different user and spawn_rate at
        different stages.

        Keyword arguments:
            stages -- A list of dicts, each representing a stage with the following keys:
                duration -- When this many seconds pass the test is advanced to the next stage #持续时间
                users -- Total user count 用户数
                spawn_rate -- Number of users to start/stop per second 每秒产生或停止的用户数
                user_classes -- 指定的任务让负载更精确的控制该任务，如果不指定就是随机的
                stop -- A boolean that can stop that test at a specific stage 要想在哪个阶段停止运行就设置该值

            stop_at_end -- Can be set to stop once all stages have run.
        """

        stages = runConf.get('stages',[])

        def tick(self):
            run_time = self.get_run_time()

            for stage in self.stages:
                if run_time < stage["duration"]:
                    try:
                        tick_data = (stage["users"], stage["spawn_rate"], stage["user_classes"])
                    except KeyError:
                        tick_data = (stage["users"], stage["spawn_rate"])
                    return tick_data

            return None

if __name__ == '__main__':
    subprocess.run(
        ["locust", "-f", "performanceengine/taskGenerate.py"])
    # test = GenerateTask().create_taskSet_class(item)
    # test = BaseStepDispose().api_perform(item.get('scenes')[0].get('steps')[0])
    # print(test)

